import torch
from torch import nn as nn


class DuelQNet(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.l1 = nn.Linear(state_dim, 256)
        self.l2 = nn.Linear(256, 256)
        self.sv_head = nn.Linear(256, 1)
        self.adv_head = nn.Linear(256, action_dim)

    def forward(self, x):
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        adv = self.adv_head(x)
        sv = self.sv_head(x)
        q_values = sv + adv - torch.mean(adv, dim=1).unsqueeze(1)
        return q_values
